from tensorflow import keras
import tensorflow as tf
from solvers.mpnn import MPNNModel


def keras2tf(f_keras, tf_dir):
  mpnn = MPNNModel(atom_dim=29, bond_dim=7)
  mpnn.compile(
    loss=keras.losses.BinaryCrossentropy(),
    optimizer=keras.optimizers.Adam(learning_rate=5e-4),
    metrics=[keras.metrics.AUC(name="AUC")],
  )
  mpnn.load_weights(f_keras)
  tf.saved_model.save(mpnn, tf_dir)

if __name__ == '__main__':
  keras2tf('weights/mpnn_bbbp_e40.h5', 'weights/mpnn_bbbp_e40')